import json
import pandas as pd
from tqdm import tqdm
from pathlib import Path
import uuid
from loguru import logger
import random
from subprocess import Popen, PIPE, CalledProcessError
import shutil

# meloTTS
# from melo.api import TTS
# import MeCab    # need run python -m unidic download
# model = TTS(language='ZH', device='cpu')
# def generate(text: str, output_path: str):
#     speaker_ids = model.hps.data.spk2id
#     model.tts_to_file(text, speaker_ids['ZH'], output_path, speed=1.0)
#     return output_path

# chatTTS
# rvcmd -w 1 assets/chtts to download assets
import ChatTTS
import torch
import torchaudio
torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.suppress_errors = True

chat = ChatTTS.Chat()
chat.load(compile=False)
def generate(text: str, output_path: str):
    wavs = chat.infer(text)
    torchaudio.save(output_path, torch.from_numpy(wavs[0]).unsqueeze(0), 24000)

if __name__ == '__main__':
    n_files_per_token = [2,3]
    df = pd.read_csv('pacenote_view.csv', index_col='id')
    inter_path = Path('/tmp/') / str(uuid.uuid4())
    inter_path.mkdir(exist_ok=True)
    logger.info(f'using intermedite path {inter_path}')
    output_path = Path('./output')
    output_path.mkdir(exist_ok=True)
    result = {} # {"filename": [list of files]}
    for idx, row in tqdm(list(df.iterrows())):
        desc = row['description']
        primary_filename = row['primary_filename']
        # split desc by '/'
        txts = desc.split('/')
        for i, txt in enumerate(txts):
            for j in range(random.randint(*n_files_per_token)):
                # create random tmp file
                try:
                    tmp_file_uuid = str(uuid.uuid4())
                    tmp_file = inter_path / f'{tmp_file_uuid}.wav'
                    generate(txt, tmp_file)
                    if primary_filename not in result:
                        result[primary_filename] = []
                    result[primary_filename].append(tmp_file_uuid)
                except:
                    pass

    # save result to file
    with open(output_path / 'result.json', 'w') as f:
        json.dump(result, f)
    
    # start ftransc command to convert wav to ogg
    cmd = f'ftransc --directory {inter_path} -f ogg'
    logger.info(f'going to run command {cmd}')
    with Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE) as p:
        for line in p.stdout:
            print(line, end='') # process line here

    logger.info(f'ftransc command finished')
    logger.info(f'arranging files...')
    # cp files to output path
    for primary_filename, files in tqdm(result.items()):
        # mkdir 
        output_path_token = output_path / primary_filename
        output_path_token.mkdir(exist_ok=True)
        for file in files:
            src = inter_path / f'{file}.ogg'
            dst = output_path_token / f'{file}.ogg'
            logger.info(f'copying {src} to {dst}')
            shutil.copy(src, dst)
    
    logger.info('done')




